# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import logging
import math
import functools
import signal
import random

from absl import flags, app

from coin.base.timestamp import get_timestamp
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str2
from coin.base.param_util import to_list
import coin.proto.account_control_pb2 as control_proto
from coin.proto.coin_request_pb2 import AccountRequestProto
from coin.strategy.executor.passive.factory import construct_mm_executor
from coin.strategy.executor.aggressive.factory import construct_agg_executor
from coin.strategy.executor.util.cancel_all_executor import (SimpleCancelAllExecutor)
from coin.strategy.mm.base.unary_multifeed_strategy_base import UnaryMultifeedStrategyBase
from coin.strategy.mm.subscription import SubscriptionRequest
from coin.strategy.mm.live_envs import get_live_env_cls
from coin.proto.coin_executor_pb2 import (PassiveExecutorConfig, AggressiveExecutorConfig)
from coin.exchange.base.order_gateway import OrderSide
from coin.strategy.mm.subscription import (gen_feed_sub_request_with_prefer)

FLAGS = flags.FLAGS


def ceil_price(price, tick_size):
  return math.ceil(price / tick_size) * tick_size


def floor_price(price, tick_size):
  return math.floor(price / tick_size) * tick_size


def _check_request(request: control_proto.ControlRequest):
  assert request.target.HasField('market_type')
  assert request.target.HasField('exchange')
  assert request.HasField('symbol')
  assert request.type == control_proto.CONTROL_APPROACH_POSITION
  approach_request = request.approach_position
  assert approach_request.HasField('target_position')
  assert approach_request.HasField('target_error_tolerance')
  assert approach_request.HasField('pass_bp')
  assert approach_request.HasField('lot_size')
  assert approach_request.HasField('stickiness')
  assert approach_request.HasField('order_update_period')


class ApproachTargetPositionStrategy(UnaryMultifeedStrategyBase):
  def __init__(self,
               control_request: control_proto.ControlRequest,
               dry_run=True,
               logger=None,
               rand_lot_size_percent=0.0):
    _check_request(control_request)
    market_type = control_request.target.market_type
    exchange = control_request.target.exchange
    api_version = control_request.target.api_version if \
                  control_request.target.HasField('api_version') else None
    symbol = control_request.symbol
    product = generate_product_from_str2(market_type=market_type,
                                         exchange=exchange,
                                         norm_product=symbol,
                                         api_version=api_version)
    execution_request = control_request.approach_position

    self._trade_sub_req = SubscriptionRequest(market_type=market_type,
                                              exchange=exchange,
                                              api_version=api_version)
    order_cls = get_live_env_cls(self._trade_sub_req)
    self._set_traits(self._trade_sub_req)

    currencies = [product.base, product.quote]
    currencies_str = [product.base.currency, product.quote.currency]
    if 'ETH' not in currencies_str:
      # hacking to add ETH as currency
      product_dummy = generate_product_from_str2(market_type=market_type,
                                                 exchange=exchange,
                                                 norm_product='ETH-USDC',
                                                 api_version=api_version)
      currencies.append(product_dummy.base)
    super().__init__(products=to_list(product), currencies=currencies)

    self._lot_size = control_request.approach_position.lot_size
    self._rand_lot_size = rand_lot_size_percent / 100.0
    self._market_type = market_type
    self._exchange = exchange
    self._api_version = api_version
    self._control_request = control_request
    self._product = product
    self._execution_request = execution_request
    self._target_position = execution_request.target_position
    self._pass_bp = execution_request.pass_bp
    self._tolerance = execution_request.target_error_tolerance
    self._prev_log_time = 0
    self._dry_run = dry_run
    self._executor = None
    self._target_reached = False
    self._strat_cancelled = False
    self._exit = False
    self._exit_attempt = 0
    self._exit_executor = None
    self._logger = logger or logging.getLogger('ApproachTargetPositionStrategy')
    self._total_fills = 0.0
    self._first_book_ts = None

  @property
  def strat_cancelled(self):
    return self._strat_cancelled

  def set_exit(self):
    self._exit = True

  def get_ask_qty_sum(self, book, num_level):
    qty_sum = 0
    for price, level in book.iteritems_ask():
      if level.qty > 0:
        qty_sum += level.qty
        num_level -= 1
        if num_level <= 0:
          break
    return qty_sum

  def get_bid_qty_sum(self, book, num_level):
    qty_sum = 0
    for price, level in book.iteritems_bid():
      if level.qty > 0:
        qty_sum += level.qty
        num_level -= 1
        if num_level <= 0:
          break
    return qty_sum

  def prepare(self):
    if self._exchange == "Uniswap":
      prefer_subreq = SubscriptionRequest(market_type=self._market_type,
                                          exchange=self._exchange,
                                          api_version=self._api_version)
    else:
      prefer_subreq = None
    feed_sub_request = gen_feed_sub_request_with_prefer(self._product, prefer_subreq)
    self.feed_subsystem.subscribe_book(feed_sub_request,
                                       functools.partial(self.on_book, self._product))

  def get_position(self):
    if self._market_type == "Spot":
      actual_position = self.order_gateway.get_balance(self._product.base.currency)
    elif self._market_type == "Futures":
      actual_position = self.order_gateway.get_position(self._product)
    else:
      raise ValueError(self._market_type)
    return actual_position

  def on_og_reset(self):
    order_gateway = self.order_gateway
    actual_position = self.get_position()

    long_position = None
    short_position = None

    if flags.FLAGS.hard_close:
      positions_proto = self.order_gateway.og_info.get_product_position_proto(self._product)
      long_position = positions_proto.long_position
      short_position = positions_proto.short_position
    elif abs(actual_position - self._target_position) <= self._tolerance:
      self._logger.info('Target position already achieved.')
      self.on_exit()
      return

    if flags.FLAGS.hard_close and long_position == 0 and short_position == 0:
      self._logger.info('%s, %s, Target position already achieved.', flags.FLAGS.account,
                        self._product)
      self.on_exit()
      return

    self._logger.info('Curr position: %s vs target position: %s' %
                      (actual_position, self._target_position))
    self._logger.info('Intializing passive executor...')
    if flags.FLAGS.exchange == 'Uniswap':
      construct_func = construct_agg_executor
      exec_proto = AggressiveExecutorConfig
    else:
      construct_func = construct_mm_executor
      exec_proto = PassiveExecutorConfig
    if flags.FLAGS.hard_close:
      self._executor = construct_func(
          self._trade_sub_req,
          exec_proto(
              symbol=self._product.symbol,
              lot_size=self._lot_size * (1 + self._rand_lot_size),
              min_pos=-self._lot_size,
              max_pos=self._lot_size * (1 + self._rand_lot_size),
              order_update_period=int(self._execution_request.order_update_period * 1e9),
              min_posting_period=int(1e8),
              max_posting_period=int(120e9),
              fill_cooldown_period=int(1e8),
          ), order_gateway, self._logger)
    else:
      min_pos = min(self._target_position, actual_position)
      max_pos = max(self._target_position, actual_position)
      self._executor = construct_func(
          self._trade_sub_req,
          exec_proto(
              symbol=self._product.symbol,
              lot_size=self._lot_size * (1 + self._rand_lot_size),
              min_pos=min_pos,
              max_pos=max_pos,
              order_update_period=int(self._execution_request.order_update_period * 1e9),
              min_posting_period=int(1e8),
              max_posting_period=int(120e9),
              fill_cooldown_period=int(1e8),
          ), order_gateway, self._logger)
    if flags.FLAGS.hard_close:
      self._executor._traits.set_hard_close()
    order_gateway.register_fill_callback(self.on_passive_fill)
    order_gateway.cancel_all()
    self._exit_executor = SimpleCancelAllExecutor()

  def _check_target_reached(self):
    if self._target_reached:
      return True
    curr_position = self.get_position()
    if flags.FLAGS.hard_close:
      positions_proto = self.order_gateway.og_info.get_product_position_proto(self._product)
      long_position = positions_proto.long_position
      short_position = positions_proto.short_position
      if long_position == 0 and short_position == 0:
        self._target_reached = True
        self._logger.info('Target position achieved.')
        self._ioloop.add_callback(self.on_exit)
    elif abs(curr_position - self._target_position) <= self._tolerance:
      self._target_reached = True
      self._logger.info('Target position achieved.')
      self._ioloop.add_callback(self.on_exit)

  def on_book(self, product, book):
    if self._exit:
      return

    if product != self._product:
      return

    if not self.order_gateway or not self.order_gateway.is_ready():
      return

    if self._first_book_ts is None:
      self._first_book_ts = book.timestamp

    if book.timestamp - self._first_book_ts > flags.FLAGS.auto_shutdown_sec * 1e9:
      self.on_exit()
      return

    tick_size = 0.5
    pass_bp = self._pass_bp
    ask_qty_sum = self.get_ask_qty_sum(book, 2)
    bid_qty_sum = self.get_bid_qty_sum(book, 2)

    ask0 = book.ask0().price
    askt = book.get_notional_asks_by_qty(bid_qty_sum)[1][-1][0]
    bid0 = book.bid0().price
    bidt = book.get_notional_bids_by_qty(ask_qty_sum)[1][-1][0]

    sell_price = min(ceil_price(bid0 * (1 + pass_bp / 10000.), tick_size), askt)
    buy_price = max(floor_price(ask0 * (1 - pass_bp / 10000.), tick_size), bidt)

    if not self._dry_run and self._executor is not None:
      if self._check_target_reached():
        return

      curr_position = self.get_position()

      if self._prev_log_time + 1e9 < book.timestamp:
        self._logger.info("ask0 %s bid0 %s sellp %s buyp %s pass_bp %s", ask0, bid0, sell_price,
                          buy_price, self._pass_bp)
        self._logger.info('Curr position: %s vs target position: %s' %
                          (curr_position, self._target_position))
        self._prev_log_time = book.timestamp
      if FLAGS.exchange == 'Uniswap':
        qty = self._lot_size
        if self._target_position >= curr_position:
          side = OrderSide.BUY
          price = ask0
        elif self._target_position <= curr_position:
          side = OrderSide.SELL
          price = bid0
        else:
          raise ValueError()
        self._executor.submit_agg_order(
            order_side=side,
            order_price=price,
            order_qty=qty * (1 + random.uniform(-self._rand_lot_size, self._rand_lot_size)))
        return

      if self._target_position >= curr_position:
        sell_price = None
        sell_qty = None
        buy_qty = self._lot_size * (1 +
                                    random.uniform(-self._rand_lot_size, self._rand_lot_size))
      elif self._target_position <= curr_position:
        buy_price = None
        buy_qty = None
        sell_qty = self._lot_size * (
            1 + random.uniform(-self._rand_lot_size, self._rand_lot_size))
      else:
        raise ValueError()

      self._executor.manage_mm_orders(sell_price=sell_price,
                                      buy_price=buy_price,
                                      sell_qty=sell_qty,
                                      buy_qty=buy_qty)

  def on_passive_fill(self, fill_price, fill_qty, order):
    print((fill_price, fill_qty, order))
    self._total_fills += fill_qty

  def signal_handler(self, sig, frame):
    self._logger.warning('Caught signal: %s', sig)
    self._strat_cancelled = True
    self._ioloop.add_callback_from_signal(self.on_exit)

  def on_exit(self):
    self.set_exit()

    self._exit_attempt += 1
    if self._exit_attempt > 5:
      self._ioloop.add_callback(self._ioloop.stop)

    finished = True
    if self._exit_executor is not None:
      ogs = to_list(self.order_gateway)
      self._logger.info('Exit callback executed.')
      finished = self._exit_executor.cancel_all_on_exit(ogs)

    if finished:
      self._ioloop.add_callback(self._ioloop.stop)
    else:
      timeout = datetime.timedelta(seconds=2)
      self._ioloop.add_timeout(timeout, self.on_exit)
    self._logger.info(f'Total abs fill qty {self._total_fills}')

  '''
  def start(self, og_config_name):
    self.run_live(
      strategy=self,
      og_config_name=og_config_name,
      on_order_gateway_reset_callback=self.on_og_reset)
    curr_position = self.order_gateway.get_position(self._product)
    success = abs(curr_position - self._target_position) <= self._tolerance
    approach_position = control_proto.ApproachPositionResponse(
        success=success, final_position=curr_position)
    control_response = control_proto.ControlResponse(
        target=self._control_request.target,
        type=control_proto.CONTROL_APPROACH_POSITION,
        approach_position=approach_position)
    return control_response
    '''


def main(argv):
  exchange = FLAGS.exchange
  api_version = FLAGS.api_version
  market_type = FLAGS.market_type
  account = FLAGS.account
  og_config_name = FLAGS.og_config_name
  symbol = FLAGS.symbol
  target_position = FLAGS.target_position
  pass_bp = FLAGS.pass_bp
  lot_size = FLAGS.lot_size
  tolerance = FLAGS.tolerance
  stickiness = FLAGS.stickiness
  order_update_period = FLAGS.order_update_period
  dry_run = FLAGS.dry_run
  rand_lot_size_percent = FLAGS.rand_lot_size_percent

  assert exchange, '--exchange must be specified.'
  assert account, '--account must be specified.'
  assert symbol, '--symbol must be specified.'
  assert lot_size, '--lot_size must be specified.'

  target = AccountRequestProto(market_type=market_type,
                               exchange=exchange,
                               api_version=api_version,
                               owner=account)
  approach_position = control_proto.ApproachPositionRequest(target_position=target_position,
                                                            pass_bp=pass_bp,
                                                            lot_size=lot_size,
                                                            target_error_tolerance=tolerance,
                                                            stickiness=stickiness,
                                                            order_update_period=order_update_period)
  control_request = control_proto.ControlRequest(target=target,
                                                 symbol=symbol,
                                                 type=control_proto.CONTROL_APPROACH_POSITION,
                                                 approach_position=approach_position)

  strategy = ApproachTargetPositionStrategy(control_request,
                                            dry_run=dry_run,
                                            rand_lot_size_percent=rand_lot_size_percent)
  signal.signal(signal.SIGTERM, strategy.signal_handler)
  signal.signal(signal.SIGINT, strategy.signal_handler)

  return strategy.run_live(strategy=strategy,
                           on_order_gateway_reset_callback=strategy.on_og_reset,
                           og_config_name=og_config_name,
                           og_params=None,
                           init_order_gateway_after=datetime.timedelta(minutes=0.1),
                           exit_after=datetime.timedelta(minutes=10),
                           fs_config=None,
                           og_log_callback=None)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)s %(asctime)s %(name)s] %(message)s')

  flags.DEFINE_boolean('dry_run', False, 'Dry run.')

  flags.DEFINE_string('exchange', None, 'Exchange name.')

  flags.DEFINE_string('api_version', None, '')

  flags.DEFINE_string('market_type', None, '')

  flags.DEFINE_string('account', None, 'Trading account name.')

  flags.DEFINE_string('og_config_name', None, 'og_config_name')

  flags.DEFINE_string('symbol', None, 'Symbol name.')

  flags.DEFINE_float('target_position', 0, 'Target position to approach.')

  flags.DEFINE_float('tolerance', 0.1, 'target_error_tolerance')

  flags.DEFINE_float('pass_bp', 5, 'Agg bp for bid ask price.')

  flags.DEFINE_float('lot_size', None, 'Lot size for each post.')

  flags.DEFINE_float(
    'stickiness',
    0,
    'If difference between price of new order and existing working order is '\
    'greater than stickness, place a new order.')

  flags.DEFINE_float('order_update_period', 2, 'Time in seconds to place a new order.')
  flags.DEFINE_boolean('hard_close', False, 'perform hard close')
  flags.DEFINE_float('rand_lot_size_percent', 0.0, 'randomize lot percent')
  flags.DEFINE_float('auto_shutdown_sec', 10.0, 'auto shutdown after x seconds')

  app.run(main)
